package cn.jly.bigdata.spark.streaming

import java.io.{BufferedReader, InputStreamReader}
import java.net.Socket

import org.apache.spark.SparkConf
import org.apache.spark.storage.StorageLevel
import org.apache.spark.streaming.dstream.ReceiverInputDStream
import org.apache.spark.streaming.{Seconds, StreamingContext}
import org.apache.spark.streaming.receiver.Receiver

import scala.util.control.Breaks

/**
 * 自定义数据源，即自定义数据采集器
 * 实现监控某个端口号，获取该端口号内容
 *
 * @author lanyangji
 * @date 2019/12/4 14:06
 */
object SparkStreaming04_CustomDataSource {

  def main(args: Array[String]): Unit = {

    val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("SparkStreaming04_CustomDataSource")
    val streamingContext = new StreamingContext(sparkConf, Seconds(5))

    // 使用自定义Receiver的Streaming
    val socketLineStream: ReceiverInputDStream[String] = streamingContext.receiverStream(new CustomReceiver("hadoop102", 9999))

    // 处理
    socketLineStream.flatMap(_.split(" ")).map((_, 1)).reduceByKey(_ + _).print()

    // 启动StreamingContext
    streamingContext.start()
    streamingContext.awaitTermination()
  }
}

// 自定义数据采集器
// 1. 继承 Receiver抽象类
// 2. 重写 onStart 和 onStop 方法
class CustomReceiver(host: String, port: Int) extends Receiver[String](StorageLevel.MEMORY_ONLY) {

  var socket: Socket = _
  // 创建一个BufferedReader用于接收端口传来的数据
  var bufferedReader: BufferedReader = _

  // 最初启动的时候，调用该方法，作用为：读数据并将数据发送给spark
  override def onStart(): Unit = {

    // 创建一个线程来读数据
    new Thread("My Socket Receiver") {

      // 重写run方法
      override def run(): Unit = {
        receive()
      }
    }.start()
  }

  /**
   * 读取数据并将数据发送给spark
   */
  def receive(): Unit = {

    // 创建一个socket
    socket = new Socket(host, port)
    // 创建一个BufferedReader用于接收端口传来的数据
    bufferedReader = new BufferedReader(new InputStreamReader(socket.getInputStream, "UTF-8"))
    // 定义一个变量，用来接收端口传来的数据一行一行的数据
    var input: String = null

    // 读取数据
    input = bufferedReader.readLine()
    // 当receiver没有关闭并且输入的数据不为空，则循环发送数据给Spark
    Breaks.breakable {
      while (!isStopped() && input != null) {

        // 如果读取的数据包含“END”，则断开接收器
        if (input.contains("END")) {
          Breaks.break()
        }

        // 将数据发送给spark
        this.store(input)

        // 继续读下一行
        input = bufferedReader.readLine()
      }
    }

    // 跳出循环后释放资源
    bufferedReader.close()
    socket.close()

    // 重启任务
    restart("CustomReceiver...restarting...")
  }

  override def onStop(): Unit = {

    // 关流
    if (bufferedReader != null) {
      bufferedReader.close()
    }

    // 释放资源
    if (socket != null) {
      socket.close()
      // 让JVM立即回收，即使可以不写，close底层有这么做了
      socket = null
    }
  }
}